package org.apache.flink.connector.scylladb.serialization;

import com.datastax.driver.core.Row;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.data.StringData;
import org.apache.flink.table.data.ArrayData;
import org.apache.flink.table.data.GenericArrayData;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.logical.ArrayType;
import static org.apache.flink.table.types.logical.LogicalTypeRoot.*;

import java.math.BigDecimal;
import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.Date;
import java.util.Set;
import java.util.UUID;

/**
 * ScyllaDB 行数据反序列化器，将 Cassandra Row 转换为 Flink RowData
 */
public class ScyllaDBRowDataDeserializer {

    private final RowType rowType;
    private final int[][] projectedFields;

    public ScyllaDBRowDataDeserializer(RowType rowType, int[][] projectedFields) {
        this.rowType = rowType;
        this.projectedFields = projectedFields;
    }

    public RowData deserialize(Row row) {
        if (projectedFields != null && projectedFields.length > 0) {
            return deserializeProjected(row);
        } else {
            return deserializeAll(row);
        }
    }

    private RowData deserializeAll(Row row) {
        GenericRowData rowData = new GenericRowData(rowType.getFieldCount());
        
        for (int i = 0; i < rowType.getFieldCount(); i++) {
            String fieldName = rowType.getFieldNames().get(i);
            LogicalType fieldType = rowType.getTypeAt(i);
            Object value = convertValue(row, fieldName, fieldType, i);
            rowData.setField(i, value);
        }
        
        return rowData;
    }

    private RowData deserializeProjected(Row row) {
        GenericRowData rowData = new GenericRowData(projectedFields.length);
        
        for (int i = 0; i < projectedFields.length; i++) {
            int fieldIndex = projectedFields[i][0];
            String fieldName = rowType.getFieldNames().get(fieldIndex);
            LogicalType fieldType = rowType.getTypeAt(fieldIndex);
            Object value = convertValue(row, fieldName, fieldType, fieldIndex);
            rowData.setField(i, value);
        }
        
        return rowData;
    }

    private Object convertValue(Row row, String fieldName, LogicalType fieldType, int fieldIndex) {
        if (row.isNull(fieldName)) {
            return null;
        }

        switch (fieldType.getTypeRoot()) {
            case CHAR:
            case VARCHAR:
                return StringData.fromString(row.getString(fieldName));
            case BOOLEAN:
                return row.getBool(fieldName);
            case TINYINT:
                return (byte) row.getInt(fieldName);
            case SMALLINT:
                return (short) row.getInt(fieldName);
            case INTEGER:
                return row.getInt(fieldName);
            case BIGINT:
                return row.getLong(fieldName);
            case FLOAT:
                return row.getFloat(fieldName);
            case DOUBLE:
                return row.getDouble(fieldName);
            case DECIMAL:
                return row.getDecimal(fieldName);
            case TIMESTAMP_WITHOUT_TIME_ZONE:
                Date timestamp = row.getTimestamp(fieldName);
                if (timestamp != null) {
                    return timestamp.toInstant().atOffset(ZoneOffset.UTC).toLocalDateTime();
                }
                return null;
            case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
                Date timestampWithTz = row.getTimestamp(fieldName);
                if (timestampWithTz != null) {
                    return timestampWithTz.toInstant();
                }
                return null;
            case BINARY:
            case VARBINARY:
                return row.getBytes(fieldName).array();
            case ARRAY:
                // 处理 Set 类型（在 ScyllaDB 中是 set<text>，在 Flink 中表示为 ARRAY 类型）
                ArrayType arrayType = (ArrayType) fieldType;
                LogicalType elementType = arrayType.getElementType();
                
                if (elementType.getTypeRoot() == VARCHAR || 
                    elementType.getTypeRoot() == CHAR) {
                    // 处理 set<text> 类型
                    Set<String> setData = row.getSet(fieldName, String.class);
                    if (setData != null) {
                        StringData[] stringArray = setData.stream()
                                .map(StringData::fromString)
                                .toArray(StringData[]::new);
                        return new GenericArrayData(stringArray);
                    }
                } else {
                    // 处理其他元素类型的 Set
                    Set<Object> setData = row.getSet(fieldName, Object.class);
                    if (setData != null) {
                        Object[] objectArray = setData.toArray();
                        return new GenericArrayData(objectArray);
                    }
                }
                return null;
            default:
                // 对于其他类型，尝试转换为字符串
                Object obj = row.getObject(fieldName);
                if (obj != null) {
                    return StringData.fromString(obj.toString());
                }
                return null;
        }
    }
} 